import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

# 初始化pro接口
pro = ts.pro_api('a5acdd993403f9e5ae13a16f2d9182b8e35d67e3beddfebd4bb43712')


def get_stock_data(ts_codes, start_date, end_date):
    data_frames = []
    for ts_code in ts_codes:
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        df = df.sort_values(by='trade_date')
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        data_frames.append(df)
    return data_frames


def simple_trading_strategy(df):
    # 简单的交易信号：金叉（短期均线穿过长期均线）作为买点，死叉作为卖点
    short_window = 5
    long_window = 20
    df['short_ma'] = df['close'].rolling(window=short_window).mean()
    df['long_ma'] = df['close'].rolling(window=long_window).mean()
    df['signal'] = 0
    df.loc[df['short_ma'] > df['long_ma'], 'signal'] = 1
    df.loc[df['short_ma'] < df['long_ma'], 'signal'] = -1
    df['position'] = df['signal'].diff()

    # 止损与止盈
    stop_loss = 0.05
    take_profit = 0.1
    positions = []
    current_position = 0
    for i in range(len(df)):
        if df['position'].iloc[i] == 1:  # 买点
            current_position = 1
            buy_price = df['close'].iloc[i]
        elif df['position'].iloc[i] == -1:  # 卖点
            current_position = 0
        elif current_position == 1:
            current_price = df['close'].iloc[i]
            if (current_price / buy_price - 1) <= -stop_loss or (current_price / buy_price - 1) >= take_profit:
                current_position = 0
        positions.append(current_position)
    df['position'] = positions

    # 计算收益率
    df['returns'] = df['close'].pct_change()
    df['strategy_returns'] = df['position'].shift(1) * df['returns']
    df['cumulative_returns'] = (1 + df['strategy_returns']).cumprod()
    return df


def visualize_results(data_frames):
    plt.figure(figsize=(15, 10))
    for df in data_frames:
        ts_code = df['ts_code'].iloc[0]
        # 绘制K线图
        plt.subplot(2, 1, 1)
        plt.plot(df['trade_date'], df['close'], label=f'{ts_code} Close Price')
        plt.scatter(df[df['position'] == 1]['trade_date'], df[df['position'] == 1]['close'],
                    marker='^', color='g', label=f'{ts_code} Buy Signal')
        plt.scatter(df[df['position'] == -1]['trade_date'], df[df['position'] == -1]['close'],
                    marker='v', color='r', label=f'{ts_code} Sell Signal')
        plt.title('Stock Price with Buy and Sell Signals')
        plt.xlabel('Date')
        plt.ylabel('Price')
        plt.legend()

        # 绘制收益率曲线
        plt.subplot(2, 1, 2)
        plt.plot(df['trade_date'], df['cumulative_returns'], label=f'{ts_code} Cumulative Returns')
        plt.title('Cumulative Returns of Trading Strategy')
        plt.xlabel('Date')
        plt.ylabel('Cumulative Returns')
        plt.legend()
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    ts_codes = ["688128.SH", "600278.SH"]  # 多支股票代码
    start_date = '20230101'
    end_date = '20250101'

    # 获取数据
    data_frames = get_stock_data(ts_codes, start_date, end_date)

    # 应用交易策略
    strategy_data_frames = []
    for df in data_frames:
        strategy_df = simple_trading_strategy(df)
        strategy_data_frames.append(strategy_df)

    # 可视化结果
    visualize_results(strategy_data_frames)